from datetime import datetime
from datetime import timedelta

from flask import Flask
from flask import jsonify
from flask import request

from flask_jwt_extended import create_access_token
from flask_jwt_extended import set_access_cookies
from flask_jwt_extended import get_jwt_identity
from flask_jwt_extended import get_jwt
from flask_jwt_extended import jwt_required
from flask_jwt_extended import JWTManager

app = Flask(__name__)

# Setup the Flask-JWT-Extended extension
app.config["JWT_SECRET_KEY"] = "x 刈刂"  # Change this!
app.config["JWT_ACCESS_TOKEN_EXPIRES"] = timedelta(days=7)
app.config["JWT_TOKEN_LOCATION"] = ["cookies"]
app.config["JWT_DECODE_LEEWAY"] = timedelta(hours=4) # 冗余4小时的过期时间
app.config["JWT_IDENTITY_CLAIM"] = 'id'
app.config["JWT_ACCESS_COOKIE_NAME"] = 'token'
app.config["JWT_COOKIE_CSRF_PROTECT"] = False
jwt = JWTManager(app)


# Using an `after_request` callback, we refresh any token that is within 30
# minutes of expiring. Change the timedeltas to match the needs of your application.
@app.after_request
def refresh_expiring_jwts(response):
    try:
        exp_timestamp = get_jwt()["exp"]
        now_timestamp = datetime.timestamp(datetime.now())
        # if target_timestamp > exp_timestamp:
        if now_timestamp > exp_timestamp:
            access_token = create_access_token(identity=get_jwt_identity())
            print(access_token)
            set_access_cookies(response, access_token)
        return response
    except (RuntimeError, KeyError):
        # Case where there is not a valid JWT. Just return the original respone
        return response


# Create a route to authenticate your users and return JWTs. The
# create_access_token() function is used to actually generate the JWT.
@app.route("/login", methods=["POST"])
def login():
    username = request.json.get("username", None)
    password = request.json.get("password", None)
    if username != "admin@knownsec.com" or password != "testpasword":
        return jsonify({"msg": "Bad username or password"}), 401

    access_token = create_access_token(identity=username)
    return jsonify(access_token=access_token)


# Protect a route with jwt_required, which will kick out requests
# without a valid JWT present.
@app.route("/protected", methods=["GET"])
@jwt_required()
def protected():
    # Access the identity of the current user with get_jwt_identity
    id = get_jwt_identity()
    jwt = get_jwt()
    print(jwt)
    return jsonify(id=id, jwt=jwt), 200


if __name__ == "__main__":
    app.run(host='0.0.0.0', port=8003)